#Script for alternative subdivision patterns for GE.
#ROIs that failed to produce 8 time series were put into
#the ROIs_for_repair dictionary. A 2x4 pattern was first
#attempted, followed by 2x2x2 and single-pass 8.

import os
import mne
from mne.minimum_norm import read_inverse_operator, apply_inverse_raw

from GE_config import study_dir, gps_dir, subjects_dir, roi_dir, subdir_names

label_div_subdir = 'correctResponseOnly_39_ROI'

mne.set_log_level(False)

#ROIs_for_repair = {'GE_05':['R_postCG_2-rh', 'R_postCG_3-rh'], 'GE_06':['R_postCG_2-rh', 'R_preCG_2-rh', 'R_STG_1-rh', 'R_STG_3-rh'], 'GE_07':['R_postCG_2-rh', 'R_STG_1-rh'], 'GE_09':['R_ITG_3-rh'], 'GE_10':['L_CC_1-lh', 'R_postCG_2-rh'], 'GE_11':['R_ITG_3-rh', 'R_postCG_4-rh', 'R_preCG_2-rh'], 'GE_12':['L_CC_1-lh'], 'GE_13':['L_SFG_1-lh', 'L_SPC_1-lh', 'R_postCG_1-rh', 'R_SFG_3-rh'], 'GE_14':['L_CC_1-lh'], 'GE_20':['R_ITG_3-rh', 'R_STG_3-rh'], 'GE_22':['L_SPC_1-lh', 'R_postCG_2-rh'], 'GE_23':['L_CC_1-lh', 'R_postCG_5-rh', 'R_STG_2-rh'], 'GE_24':['R_postCG_2-rh']} Divided 2x4
#ROIs_for_repair = {'GE_06':['R_postCG_2-rh'], 'GE_12':['L_CC_1-lh'], 'GE_13':['L_SPC_1-lh', 'L_postCG_1-lh'], 'GE_23':['L_CC_1-lh']} Divided 2x2x2
ROIs_for_repair = {'GE_13':['R_ITG_2-rh']} # Divided 4x2, will include extra for integration

for subj in ROIs_for_repair:
	labels_fname_this_subject = [ roi_dir + subj + '/' + roi + '.label' for roi in ROIs_for_repair[subj] ]
	labels = [mne.read_label(label) for label in labels_fname_this_subject]

	parts_a, parts_b, parts_c = 2, 4, 2

	for label in labels:
		labeldir = study_dir + '/' + subj + '/' + subdir_names['labels'] + '/' + label_div_subdir + '/' + label.name

		# Move the existing subdivisions to a new subdirectory
		old_labeldir = labeldir + '/original_divisions'
		if not os.path.exists(old_labeldir):
			os.mkdir(old_labeldir)
		for f in os.listdir(labeldir):
			if os.path.isfile(labeldir + '/' + f):
				os.rename(labeldir+'/'+f, old_labeldir+'/'+f)

		parts = parts_a
		label_div_a = mne.split_label(label, parts=parts, subject=subj, subjects_dir=subjects_dir, freesurfer=True)

		for sub_label in label_div_a:
			parts = parts_b
			label_div_b = mne.split_label(sub_label, parts=parts, subject=subj, subjects_dir=subjects_dir, freesurfer=True)
			for new_label in label_div_b:
				parts = parts_c
				#label_div_c = mne.split_label(subsub_label, parts=parts, subject=subj, subjects_dir=subjects_dir, freesurfer=True)
				#for new_label in label_div_c:
				try:
					#Try to make an activation timeseries from the subdivided label
					fname_inv = study_dir + '/' + subj + '/rawdata/' + subj + '_meg_eeg-inv.fif'
					fname_raw = study_dir + '/' + subj + '/rawdata/' + subj + '_B1_filtered_raw.fif'
					raw = mne.io.read_raw_fif(fname_raw) 
					inv = read_inverse_operator(fname_inv)
					snr = 1.0
					lambda2 = 1.0 / snr ** 2
					method = "dSPM"
					stc = apply_inverse_raw(raw, inv, lambda2, method, new_label, 0, 100)    
					#If it worked and there is data, go ahead and save the label. Otherwise, print a warning to the terminal
					if not stc.data.size == 0:
						fname_div = labeldir + '/' + new_label.name + '.label'
    	              		#print(fname_div)
						new_label.save(fname_div)
					else:
						print(subj, new_label.name, " does not contain any source vertices")

				except:
					print("Problem with subject: ", subj)
					print(new_label.name, " failed")
					continue


				


